import os
import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import matplotlib.pyplot as plt
from utils.image import resize_image
from retrieval import search_by_text, search_by_img

class ImageSearchGUI:
    def __init__(self):
        self.window = tk.Tk()
        self.window.title("图片检索")
        self.window.minsize(400, 300)
        
        # 创建文字搜索的输入框和按钮
        self.text_entry = tk.Entry(self.window)
        self.text_entry.pack()
        self.text_search_button = tk.Button(self.window, text="文字搜索", command=self.text_search)
        self.text_search_button.pack()
        self.clip_weight_text_entry = tk.Entry(self.window)
        self.clip_weight_text_entry.pack()
        self.color_weight_text_entry = tk.Entry(self.window)
        self.color_weight_text_entry.pack()
        self.image_search_button = tk.Button(self.window, text="以图搜图", command=self.image_search)
        self.image_search_button.pack()
        
        # 创建以图搜图的拖拽区域
        self.drop_image = None
        self.drop_label = tk.Label(self.window, text="拖拽图片文件到此处")
        self.drop_label.pack()
        self.drop_label.bind("<Button-1>", self.open_file_dialog)
        self.window.bind("<B1-Motion>", self.drag_motion)
        
    # 文字搜索按钮回调
    def text_search(self):
        query = self.text_entry.get()
        if query == "":
            return
        results = search_by_text(query, topK=8)
        self.display_text_results(results, query)
        
    def open_file_dialog(self, event):
        file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg;*.jpeg;*.png")])
        self.drop_label.configure(text=file_path)
        if os.path.exists(file_path):
            image = Image.open(file_path)
            image = resize_image(image, 200)  # 调整图片大小
            photo = ImageTk.PhotoImage(image)
            self.drop_image = photo
            self.drop_label.configure(image=photo)
            self.drop_label.image = photo  # 需要保持对PhotoImage对象的引用，否则图片不会显示
        
    def drag_motion(self, event):
        self.drop_label.configure(bg="gray")
    
    # 以图搜图按钮回调
    def image_search(self):
        self.drop_label.configure(bg="white")
        file_path = self.drop_label.cget("text")
        if os.path.exists(file_path):
            img = Image.open(file_path)
            clip_weight = float(self.clip_weight_text_entry.get())
            color_weight = float(self.color_weight_text_entry.get())
            results = search_by_img(img, weights=[clip_weight, color_weight], topK=8)
            self.display_color_results(results, file_path, [clip_weight, color_weight])
        
    def display_text_results(self, results, info):
        # 创建一个包含8个子图的图形
        fig, axes = plt.subplots(3, 4, figsize=(12, 6))
        query_ax = axes[0, 0]
        query_ax.set_title(info)
        query_ax.axis('off')

        axes[0, 1].axis('off')
        axes[0, 2].axis('off')
        axes[0, 3].axis('off')
        # 遍历图片文件和标签，并在每个子图上绘制图片和标签
        for i, result in enumerate(results[:8]):
            image = plt.imread(result[0])
            ax = axes[i // 4 + 1, i % 4]
            ax.imshow(image)
            ax.set_title(f"similarity: {result[1]}")
            ax.axis('off')

        # 调整子图之间的间距
        plt.tight_layout()
        # 显示图形
        plt.show()

    def display_color_results(self, results, query_path, weights):
        # 创建一个包含8个子图的图形
        fig, axes = plt.subplots(3, 4, figsize=(12, 6))
        query_image = plt.imread(query_path)
        query_ax = axes[0, 0]
        query_ax.imshow(query_image)
        query_ax.set_title(f"query image(clip: {weights[0]}, color: {weights[1]})")
        query_ax.axis('off')

        axes[0, 1].axis('off')
        axes[0, 2].axis('off')
        axes[0, 3].axis('off')
        # 遍历图片文件和标签，并在每个子图上绘制图片和标签
        for i, result in enumerate(results[:8]):
            image = plt.imread(result[0])
            ax = axes[i // 4 + 1, i % 4]
            ax.imshow(image)
            ax.set_title(f"similarity: {result[1]}")
            ax.axis('off')

        # 调整子图之间的间距
        plt.tight_layout()
        # 显示图形
        plt.show()
        
    def run(self):
        self.window.mainloop()

# 创建并运行图形界面
gui = ImageSearchGUI()
gui.run()